package rabbit.flt.plugins.reactor.http.plugin.proxy;

import org.springframework.cglib.proxy.InvocationHandler;
import org.springframework.cglib.proxy.Proxy;
import rabbit.flt.common.trace.TraceData;
import rabbit.flt.common.utils.StringUtils;
import reactor.core.publisher.Flux;
import reactor.netty.ByteBufFlux;
import reactor.netty.NettyInbound;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.concurrent.locks.ReentrantLock;

public class NettyInboundProxy implements InvocationHandler {

    private NettyInbound realObj;

    private ByteBufFlux byteBufFlux;

    private TraceData traceData;

    private ReentrantLock lock = new ReentrantLock();

    public static NettyInbound proxy(NettyInbound realObj, TraceData traceData) {
        NettyInboundProxy h = new NettyInboundProxy();
        h.realObj = realObj;
        h.traceData = traceData;
        return (NettyInbound) Proxy.newProxyInstance(NettyInboundProxy.class.getClassLoader(),
                new Class[]{NettyInbound.class}, h);
    }

    @Override
    public Object invoke(Object o, Method method, Object[] args) throws InvocationTargetException, IllegalAccessException {
        if ("receive".equals(method.getName())) {
            if (null == byteBufFlux) {
                initByteBufFlux(method, args);
            }
            return byteBufFlux;
        }
        return method.invoke(realObj, args);
    }

    private void initByteBufFlux(Method method, Object[] args) throws IllegalAccessException, InvocationTargetException {
        try {
            lock.lock();
            if (null == byteBufFlux) {
                ByteBufFlux bufFlux = (ByteBufFlux) method.invoke(realObj, args);
                Flux<String> flux = Flux.defer(() -> {
                    StringBuilder sb = new StringBuilder();
                    return bufFlux.asString().map(s -> {
                        sb.append(s);
                        return s;
                    }).doFinally(f -> {
                        String body = StringUtils.truncate(sb.toString(), 1024);
                        traceData.getHttpResponse().setBody(body);
                    });
                });
                byteBufFlux = ByteBufFlux.fromString(flux);
            }
        } finally {
            lock.unlock();
        }
    }
}
